// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#include <pollux/exec/aggregate_info.h>
#include <pollux/exec/aggregate.h>
#include <pollux/exec/operator.h>
#include <pollux/expression/expr.h>

namespace kumo::pollux::exec {
    namespace {
        std::vector<core::LambdaTypedExprPtr> extractLambdaInputs(
            const core::AggregationNode::Aggregate &aggregate) {
            std::vector<core::LambdaTypedExprPtr> lambdas;
            for (const auto &arg: aggregate.call->inputs()) {
                if (auto lambda =
                        std::dynamic_pointer_cast<const core::LambdaTypedExpr>(arg)) {
                    lambdas.push_back(lambda);
                }
            }

            return lambdas;
        }
    } // namespace

    std::vector<AggregateInfo> toAggregateInfo(
        const core::AggregationNode &aggregationNode,
        const OperatorCtx &operatorCtx,
        uint32_t numKeys,
        std::shared_ptr<core::ExpressionEvaluator> &expressionEvaluator,
        bool isStreaming) {
        const auto numAggregates = aggregationNode.aggregates().size();
        std::vector<AggregateInfo> aggregates;
        aggregates.reserve(numAggregates);

        const auto &inputType = aggregationNode.sources()[0]->outputType();
        const auto &outputType = aggregationNode.outputType();
        const auto step = aggregationNode.step();

        for (auto i = 0; i < numAggregates; i++) {
            const auto &aggregate = aggregationNode.aggregates()[i];
            AggregateInfo info;
            // Populate input.
            auto &channels = info.inputs;
            auto &constants = info.constantInputs;
            for (const auto &arg: aggregate.call->inputs()) {
                if (auto field =
                        dynamic_cast<const core::FieldAccessTypedExpr *>(arg.get())) {
                    channels.push_back(inputType->getChildIdx(field->name()));
                    constants.push_back(nullptr);
                } else if (
                    auto constant =
                            dynamic_cast<const core::ConstantTypedExpr *>(arg.get())) {
                    channels.push_back(kConstantChannel);
                    constants.push_back(constant->toConstantVector(operatorCtx.pool()));
                } else if (
                    auto lambda = dynamic_cast<const core::LambdaTypedExpr *>(arg.get())) {
                    for (const auto &name: lambda->signature()->names()) {
                        if (auto captureIndex = inputType->getChildIdxIfExists(name)) {
                            channels.push_back(captureIndex.value());
                            constants.push_back(nullptr);
                        }
                    }
                } else {
                    POLLUX_FAIL(
                        "Expression must be field access, constant, or lambda: {}",
                        arg->toString());
                }
            }

            info.distinct = aggregate.distinct;
            info.intermediateType = Aggregate::intermediateType(
                aggregate.call->name(), aggregate.rawInputTypes);

            // Setup aggregation mask: convert the Variable Reference name to the
            // channel (projection) index, if there is a mask.
            if (const auto &mask = aggregate.mask) {
                info.mask = inputType->as_row().getChildIdx(mask->name());
            } else {
                info.mask = std::nullopt;
            }

            auto index = numKeys + i;
            const auto &aggResultType = outputType->childAt(index);
            info.function = Aggregate::create(
                aggregate.call->name(),
                isPartialOutput(step)
                    ? core::AggregationNode::Step::kPartial
                    : core::AggregationNode::Step::kSingle,
                aggregate.rawInputTypes,
                aggResultType,
                operatorCtx.driverCtx()->queryConfig());

            auto lambdas = extractLambdaInputs(aggregate);
            if (!lambdas.empty()) {
                if (expressionEvaluator == nullptr) {
                    expressionEvaluator = std::make_shared<SimpleExpressionEvaluator>(
                        operatorCtx.execCtx()->queryCtx(), operatorCtx.execCtx()->pool());
                }
                info.function->setLambdaExpressions(lambdas, expressionEvaluator);
            }

            // Ignore sorting properties if aggregate function is not sensitive to the
            // order of inputs.
            auto *entry = getAggregateFunctionEntry(aggregate.call->name());
            const auto &metadata = entry->metadata;
            if (metadata.orderSensitive) {
                // Sorting keys and orders.
                const auto numSortingKeys = aggregate.sortingKeys.size();
                POLLUX_CHECK_EQ(numSortingKeys, aggregate.sortingOrders.size());
                info.sortingOrders = aggregate.sortingOrders;
                info.sortingKeys.reserve(numSortingKeys);
                for (const auto &key: aggregate.sortingKeys) {
                    info.sortingKeys.push_back(exprToChannel(key.get(), inputType));
                }
            }

            info.output = index;
            aggregates.emplace_back(std::move(info));
        }
        return aggregates;
    }

    std::vector<std::optional<column_index_t> > extractMaskChannels(
        const std::vector<AggregateInfo> &aggregates) {
        std::vector<std::optional<column_index_t> > masks;
        masks.reserve(aggregates.size());
        for (const auto &aggregate: aggregates) {
            masks.push_back(aggregate.mask);
        }
        return masks;
    }
} // namespace kumo::pollux::exec
